import rosbag
import numpy as np
import sensor_msgs.point_cloud2 as pc2
import open3d as o3d
import cv2
from cv_bridge import CvBridge
from tictoc import TicToc
import csv
import os
import argparse
from utils import topic2path, get_ms
from my_io import write_points


def main(args):
    assert os.path.exists(args.rosbag_dir)
    bags = os.listdir(args.rosbag_dir)
    bags.sort(key=lambda x: x)
    for bag in bags:
        args.rosbag_path = os.path.join(args.rosbag_dir, bag)
        if os.path.isfile(args.rosbag_path):
            dir_name = str.split(bag, '.')[0]
            shuffix = str.split(bag, '.')[1]
            if shuffix == "bag":
                main_worker(os.path.join(args.data_path, dir_name), args)


def main_worker(dir_name, args):
    print("开始解析<{}>......".format(args.rosbag_path))

    bag = rosbag.Bag(args.rosbag_path)
    #
    # lidar_topics = ["/pointcloud_lidar1", "/pointcloud_lidar2", "/pointcloud_lidar3", "/pointcloud_lidar4"]
    # radar_topics = ["/pointcloud_radar3"]
    # old_radar_topics = {"/radar_det_list": "/pointcloud_radar3"}
    # camera_topics = ["/camera71/compressed", "/camera72/compressed", "/camera73/compressed", "/camera74/compressed",
    #                  "/camera75/compressed", "/camera76/compressed",
    #                  "/camera77/compressed", "/camera78/compressed", "/camera79/compressed", "/camera80/compressed",
    #                  "/camera81/compressed", "/camera82/compressed"]

    lidar_topics = ["/pointcloud_lidar1", "/pointcloud_lidar2", "/pointcloud_lidar3"]
    radar_topics = ["/pointcloud_radar3"]
    old_radar_topics = {}
    camera_topics = []
    camera_ori_topics = ["/camera71", "/camera72", "/camera73", "/camera74",
                     "/camera75", "/camera76",
                      "/camera79", "/camera80",
                     "/camera82"]

    localization_topics = ["/localization_result"]
    print("保存在<{}>".format(dir_name))

    # 新建目录
    new_dir = []
    ffarther_dir = []
    farther_dir = []
    child_dir = []
    ffarther_dir.append(dir_name)
    new_dir = new_dir + ffarther_dir
    farther_dir.append(os.path.join(dir_name, "lidar"))
    farther_dir.append(os.path.join(dir_name, "radar"))  # 毫米波所以直接保存为bin:因为1)不直接标注 2)pcd不能保留r_speed和rcs
    farther_dir.append(os.path.join(dir_name, "camera"))
    farther_dir.append(os.path.join(dir_name, "localization"))
    farther_dir.append(os.path.join(dir_name, "calib"))
    new_dir = new_dir + farther_dir
    for topic in lidar_topics:
        child_dir.append(os.path.join(dir_name, "lidar", topic2path(topic)))
    for topic in radar_topics:
        child_dir.append(os.path.join(dir_name, "radar", topic2path(topic)))
    for topic in camera_topics:
        child_dir.append(os.path.join(dir_name, "camera", topic2path(topic)))
    for topic in camera_ori_topics:
        child_dir.append(os.path.join(dir_name, "camera", topic2path(topic)))
    new_dir = new_dir + child_dir
    for dir in new_dir:
        if not os.path.exists(dir):
            print("新建目录<{}>".format(dir))
            os.makedirs(dir, exist_ok=True)

    # 创建calib.yml
    fs = cv2.FileStorage(os.path.join(dir_name, 'calib', 'sensor.yml'), cv2.FileStorage_WRITE)
    matrix = np.eye(4)  # 4*4的旋转平移矩阵
    for topic in lidar_topics + radar_topics:
        sensor_id = topic[-6:]
        fs.write(sensor_id, matrix)
    fs.release()
    print("新建标定文件 in {}".format(os.path.join(dir_name, 'calib', 'sensor.yml')))

    # 创建或打开文件
    csvfile = open(os.path.join(dir_name, "localization", 'localization.csv'), mode='w', newline='')
    # 标题列表
    fieldnames = ['stamp(ms)', 'latitude', 'longitude', 'altitude', 'mercator_x', 'mercator_y', 'mercator_z',
                  'vel_x', 'vel_y', 'vel_z', 'accel_x', 'accel_y', 'accel_z',
                  'roll', 'pitch', 'yaw', 'angular_x', 'angular_y', 'angular_z']
    # 创建 DictWriter 对象
    write = csv.DictWriter(csvfile, fieldnames=fieldnames)
    # 写入表头
    write.writeheader()

    # 统计耗时
    cost = TicToc("ros解析")
    for topic, msg, t in bag:
        if topic in lidar_topics:
            lidar = pc2.read_points(msg)
            points = np.array(list(lidar))
            # # 创建一个PointCloud对象
            # pcd = o3d.geometry.PointCloud()
            # # 将随机数转换成PointCloud点数据
            # pcd.points = o3d.utility.Vector3dVector(points[:, :3])
            # # 将PointCloud点数据保存成pcd文件，格式为assii文本格式
            # o3d.io.write_point_cloud(os.path.join(dir_name,'lidar', topic2path(topic), str(get_ms(msg)) + ".pcd"), pcd,
            #                          write_ascii=True)
            assert points.shape[1] >= 4
            write_points(points[:, :4], os.path.join(dir_name, 'lidar', topic2path(topic), str(get_ms(msg)) + ".bin"))

        elif topic in radar_topics:
            radar = pc2.read_points(msg)
            points = np.array(list(radar))
            points = points[:, :5]  # x y z r_speed rcs
            file_path = os.path.join(dir_name, 'radar', topic2path(topic), str(get_ms(msg)) + ".bin")
            write_points(points, file_path)
            # # 创建一个PointCloud对象
            # pcd = o3d.geometry.PointCloud()
            # # 将随机数转换成PointCloud点数据
            # pcd.points = o3d.utility.Vector3dVector(points[:, :3])# x y z r_speed rcs
            # # 将PointCloud点数据保存成pcd文件，格式为assii文本格式
            # o3d.io.write_point_cloud(os.path.join(dir_name,'radar', topic2path(topic), str(get_ms(msg)) + ".pcd"), pcd,
            #                          write_ascii=True)

        elif topic in old_radar_topics:
            points_list = []
            for det in msg.detection_array:
                points_list.append([det.f_x, det.f_y, det.f_z, det.f_RangeRate, det.s_RCS])
            points = np.array(points_list)
            file_path = os.path.join(dir_name, 'radar', topic2path(old_radar_topics[topic]),
                                     str(get_ms(msg.detection_array[0])) + ".bin")
            write_points(points, file_path)

        elif topic in camera_topics:
            bridge = CvBridge()
            cv_image = bridge.compressed_imgmsg_to_cv2(msg, "bgr8")
            cv2.imwrite(os.path.join(dir_name, "camera", topic2path(topic), str(get_ms(msg)) + ".jpg"), cv_image)
            # cv2.imshow("Image window", cv_image)
            # cv2.waitKey(3)
        elif topic in camera_ori_topics:
            bridge = CvBridge()
            cv_image = bridge.imgmsg_to_cv2(msg, "bgr8")
            cv2.imwrite(os.path.join(dir_name, "camera", topic2path(topic), str(get_ms(msg)) + ".jpg"), cv_image)

        elif topic in localization_topics:
            # +'\t'用于取消科学计数法
            write.writerow({'stamp(ms)': str(get_ms(msg)) + '\t', 'latitude': msg.latitude, 'longitude': msg.longitude,
                            'altitude': msg.altitude, 'mercator_x': msg.mercator_x, 'mercator_y': msg.mercator_y,
                            'mercator_z': msg.mercator_z, 'vel_x': msg.vel_x, 'vel_y': msg.vel_y, 'vel_z': msg.vel_z,
                            'accel_x': msg.accel_x,
                            'accel_y': msg.accel_y, 'accel_z': msg.accel_z, 'roll': msg.roll, 'pitch': msg.pitch,
                            'yaw': msg.yaw,
                            'angular_x': msg.angular_x, 'angular_y': msg.angular_y,
                            'angular_z': msg.angular_z})

    for dir in child_dir:
        print("文件夹:<{}> 包含文件<{}>个".format(dir.split('/')[-1], len(os.listdir(dir))))

    print("---------------------------------------------------------")
    print("{}   解析完毕".format(args.rosbag_path))
    cost.toc()
    print("---------------------------------------------------------")
    bag.close()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Configuration Parameters')
    parser.add_argument('--rosbag-dir',
                        default='/media/shn/SHN/temp/rosbag',
                        help='rosbag')
    parser.add_argument('--data-path',
                        default='/home/shn/Project/parse_result')
    args = parser.parse_args()

    main(args)
